import numpy.matlib as np
from math import sqrt
from math import pow
from math import fabs
from scipy.stats import mstats

def segregatingSites(SFS):
	S = 1.0 * SFS.sum()
	return S

def averageSegregatingSites( SFSs ):
	sum = 0.0
	for s in SFSs:
		sum += 1.0 * s.sum()
	
	avg = s / len(SFSs)
	return avg

def singletons(SFS):
	s = 1.0 * SFS[1]
	return s

def heterozygosity(SFS):
	n = SFS.shape[0]
	sum = 0.0
	for i in range(1, n):
		sum += i * (n - i) * SFS[i]
	
	pi = 2.0 * sum / (n * (n - 1.0)) 
	return pi

def operationOnAllSFS(linOp, SFSs):
	ss = []
	for sfs in SFSs:
		s = operationOnSFS( linOp, sfs )
		ss.append(s)
	return ss
	
def operationOnSFS(linOp, SFS):
	dp = linOp * SFS
	ss = dp.sum()
	return ss

def getOp_SegregatingSites(n):
	op = np.ravel(np.ones(n))
	op[0] = 0.0
	return op

def getOp_Singletons(n):
	op = np.ravel(np.zeros(n))
	op[1] = 1.0
	return op

def getOp_Heterozygosity(n):
	op = np.ravel(np.zeros(n))
	
	for i in range(1, n):
		op[i] = 2.0 * i * (n - i) / (n * (n - 1.0))
	
	return op

def getOp_Theta_Pi(n):
	op = getOp_Heterozygosity(n)
	return op

def getOp_Theta_H(n):
	op = np.ravel(np.zeros(n))
	
	for i in range(1, n):
		op[i] = 2.0 * i * i / (n * (n - 1.0))
	
	return op

def getOp_Theta_Singletons(n):
	op = getOp_Singletons(n)
	op *= (n - 1.0) / n
	return op

def getOp_Theta_SegSites(n):
	op = getOp_SegregatingSites(n)
	a = a_n(n)
	op /= a
	return op

def getOp_TajimasD(n):
	pi = getOp_Theta_Pi(n)
	SS = getOp_Theta_SegSites(n)
	op = pi - SS
	return op

def getOp_FuLiD(n):
	SS = getOp_SegregatingSites(n)
	S = getOp_Singletons(n)
	
	a = a_n(n)
	SING = a*S
	op = SS - SING
	
	return op

def getOp_FuLiF(n):
	pi = getOp_Heterozygosity(n)
	S = getOp_Singletons(n)
	op = pi - S
	
	return op

def getOp_FayWuH(n):
	h = getOp_Theta_H(n)
	pi = getOp_Theta_Pi(n)
	op = pi - h
	return op
	
	
def a_n(n):
	sum = 0.0
	for i in range(1, n):
		sum += 1.0 / i
	
	return sum
		
def b_n(n):
	sum = 0.0
	for i in range(1, n):
		sum += 1.0 / (i * i)
	
	return sum

def c_n(n):
	if n < 3:
		return 1.0

	a = a_n(n)
	c = n*a - 2.0*(n-1.0)
	c *= 2.0
	c /= ((n-1.0)*(n-2.0))

	return c

def normalizeToFirst( point ):
	ratio = 1.0 / point[1]
	normPoint = point * ratio
	return normPoint

def middlePvalue( data, point ):
	m = mstats.tmean(data)
	dist = fabs( point - m )
	count = 0
	for d in data:
		if fabs( d - m ) > dist:
			count += 1
	
	return ( 1.0 * count / len(data) )




class SummaryStat:
	def calculateNumerator(self, SFS):
		return 0.0
	
	def calculate(self, SFS):
		num = self.calculateNumerator(SFS)
		
		s = segregatingSites(SFS)
		n = SFS.shape[0]
		
		den = self.variance(n, s)
		D = 1.0 * num / sqrt(den)
		
		return D

	def variance(self, n, s):
		return 1.0
	
	def calculateAll(self, SFSs):
		Ds = []
		for s in SFSs:
			Ds.append( self.calculate(s) )

		return Ds
	
	def calculateAllNumerator(self, SFSs):
		Ds = []
		for s in SFSs:
			Ds.append( self.calculateNumerator(s) )

		return Ds
	
	def calculateAllVariance(self, SFSs):
		Ds = []
		for sfs in SFSs:
			s = segregatingSites(sfs)
			n = sfs.shape[0]
		
			Ds.append( self.variance(n, s) )

		return Ds


class TajimasD(SummaryStat):

	def calculateNumerator(self, SFS):
		pi = heterozygosity(SFS)
		s = segregatingSites(SFS)
		
		n = SFS.shape[0]
		a = a_n(n)
		num = 1.0 * (pi - (s / a))
		
		return num

	def variance(self, n, s):
		e1n = self.e1_n(n)
		e2n = self.e2_n(n)
		
		v = 1.0 * (e1n * s) + (e2n * s * (s - 1.0))
		
		return v

	def e1_n(self, n):
		a = a_n(n)

		e = 1.0 * (n + 1.0) / (3.0 * (n - 1.0))
		e -= 1.0 / a
		e /= a

		return e

	def e2_n(self, n):
		a = a_n(n)
		b = b_n(n)

		e = 2.0 * (n * (n + 1.0) + 3.0) / (9.0 * n * (n - 1.0))
		e -= 1.0 * (n + 2.0) / (n * a)
		e += 1.0 * b / (a * a)
		e /= 1.0 * a * a + b

		return e

class FulLiD(SummaryStat):

	def calculateNumerator(self, SFS):
		sing = singletons(SFS)
		ss = segregatingSites(SFS)
		
		n = SFS.shape[0]
		a = a_n(n)
		num = 1.0 * ss - a*sing
		
		return num

	def variance(self, n, s):
		a = a_n( n )
		b = b_n( n )
		c = c_n( n )

		vd = 1.0 + ( (a*a)/(b + (a*a)) * (c - (n+1.0)/(n-1.0)) )
		ud = a - 1.0 - vd

		v = ud*s + vd*(s*s);
		
		return v

class FulLiF(SummaryStat):

	def calculateNumerator(self, SFS):
		pi = heterozygosity(SFS)
		sing = singletons(SFS)
		
		n = SFS.shape[0]
		num = 1.0 * pi - sing
		
		return num

	def variance(self, n, s):

		a = a_n( n )
		an1 = a_n( n+1 )
		b = b_n( n )
		c = c_n( n )

		v = c
		v += (2.0 * (n*n + n + 3.0) / (9.0 * n * (n-1.0)))
		v -= (2.0/(n-1.0))
		v /= (a*a + b)

		u = 1.0
		u += ((n+1.0)/(3.0*(n-1.0)))
		u -= (4.0*(n+1.0)/(pow(n-1.0, 2.0)) * (an1 - (2.0*n)/(n+1.0)))
		u /= a
		u -= v
		
		var = u*s + v*(s*s)

		return var
	
